from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QPushButton, QTableWidget, QTableWidgetItem,
                           QInputDialog, QMessageBox, QLabel, QComboBox)
from PyQt5.QtCore import Qt
import sys
from snapshot import EPICSSnapshot
from datetime import datetime
from config_loader import get_cavity_names

class SnapshotGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.snapshot_manager = EPICSSnapshot()
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle('EPICS 快照管理工具')
        self.setGeometry(100, 100, 800, 600)

        # 创建中央部件和主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 创建按钮布局
        button_layout = QHBoxLayout()

        # 创建按钮
        self.save_button = QPushButton('保存快照', self)
        self.load_button = QPushButton('加载快照', self)
        self.refresh_button = QPushButton('刷新列表', self)

        # 添加按钮到布局
        button_layout.addWidget(self.save_button)
        button_layout.addWidget(self.load_button)
        button_layout.addWidget(self.refresh_button)

        # 创建表格
        self.table = QTableWidget()
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderLabels(['文件名', '时间戳', '描述'])
        self.table.setSelectionBehavior(QTableWidget.SelectRows)
        self.table.setEditTriggers(QTableWidget.NoEditTriggers)

        # 设置表格列宽
        self.table.setColumnWidth(0, 300)
        self.table.setColumnWidth(1, 200)
        self.table.setColumnWidth(2, 250)

        # 添加状态标签
        self.status_label = QLabel('')
        
        # 添加腔体选择下拉框
        cavity_layout = QHBoxLayout()
        self.cavity_label = QLabel('选择腔体:', self)
        self.cavity_combo = QComboBox(self)
        self.cavity_combo.addItem('所有腔体')
        self.cavity_combo.addItems(get_cavity_names())
        
        cavity_layout.addWidget(self.cavity_label)
        cavity_layout.addWidget(self.cavity_combo)
        cavity_layout.addStretch()

        # 将所有部件添加到主布局
        layout.addLayout(cavity_layout)
        layout.addLayout(button_layout)
        layout.addWidget(self.table)
        layout.addWidget(self.status_label)

        # 连接信号和槽
        self.save_button.clicked.connect(self.save_snapshot)
        self.load_button.clicked.connect(self.load_snapshot)
        self.refresh_button.clicked.connect(self.refresh_snapshots)

        # 初始化显示快照列表
        self.refresh_snapshots()

    def save_snapshot(self):
        description, ok = QInputDialog.getText(
            self, '保存快照', '请输入快照描述:'
        )
        if ok:
            cavity = None if self.cavity_combo.currentText() == '所有腔体' else self.cavity_combo.currentText()
            success, result = self.snapshot_manager.take_snapshot(cavity, description)
            if success:
                self.status_label.setText(f'快照已保存到: {result}')
                self.refresh_snapshots()
            else:
                QMessageBox.critical(self, '错误', f'保存快照失败：\n{result}')

    def load_snapshot(self):
        selected_rows = self.table.selectionModel().selectedRows()
        if not selected_rows:
            QMessageBox.warning(self, '警告', '请先选择一个快照')
            return

        row = selected_rows[0].row()
        filename = self.table.item(row, 0).text()
        cavity = None if self.cavity_combo.currentText() == '所有腔体' else self.cavity_combo.currentText()
        filepath = f"{self.snapshot_manager.snapshot_dir}/{filename}"

        reply = QMessageBox.question(
            self, '确认加载',
            f'确定要加载选中的快照到{self.cavity_combo.currentText()}吗？\n{filename}',
            QMessageBox.Yes | QMessageBox.No
        )

        if reply == QMessageBox.Yes:
            success, message = self.snapshot_manager.load_snapshot(filepath, cavity)
            if success:
                self.status_label.setText(message)
            else:
                QMessageBox.critical(self, '错误', f'加载快照失败：\n{message}')

    def refresh_snapshots(self):
        cavity = None if self.cavity_combo.currentText() == '所有腔体' else self.cavity_combo.currentText()
        snapshots = self.snapshot_manager.list_snapshots(cavity)
        self.table.setRowCount(len(snapshots))
        
        for row, snapshot in enumerate(snapshots):
            self.table.setItem(row, 0, QTableWidgetItem(snapshot['filename']))
            self.table.setItem(row, 1, QTableWidgetItem(snapshot['timestamp']))
            self.table.setItem(row, 2, QTableWidgetItem(
                f"{snapshot['description']} ({snapshot['cavity']})"
            ))

        self.status_label.setText(f'共找到 {len(snapshots)} 个快照')

def main():
    app = QApplication(sys.argv)
    window = SnapshotGUI()
    window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main() 